//
//  RTMPConnection.swift
//  replaykit
//
//  Created by zhang on 2025/7/31.
//

import Foundation
import Network

class RTMPConnection: NSObject {
    
    // MARK: - Properties
    private var isConnected = false
    private var isPublishing = false
    
    // RTMP配置
    private let rtmpURL: String
    private let streamName: String
    
    // 网络连接
    private var connection: NWConnection?
    private var host: String = "192.168.1.25"
    private var port: UInt16 = 1935
    
    // 连接超时
    private let connectionTimeout: TimeInterval = 30.0
    private var connectionTimer: Timer?
    
    // RTMP状态
    private var chunkSize: Int = 128
    private var windowAckSize: Int = 2500000
    private var peerBandwidth: Int = 2500000
    private var limitType: UInt8 = 2
    
    // 回调
    var onConnectionStatusChanged: ((Bool) -> Void)?
    var onPublishStatusChanged: ((Bool) -> Void)?
    var onError: ((String) -> Void)?
    
    // MARK: - Initialization
    init(url: String, streamName: String) {
        self.rtmpURL = url
        self.streamName = streamName
        super.init()
        print("RTMPConnection: 初始化RTMP连接管理器")
        print("RTMPConnection: URL = \(url)")
        print("RTMPConnection: StreamName = \(streamName)")
    }
    
    deinit {
        disconnect()
    }
    
    // MARK: - Connection Management
    func connect() -> Bool {
        print("RTMPConnection: 开始连接RTMP服务器: \(rtmpURL)")
        
        // 解析RTMP URL
        guard let url = URL(string: rtmpURL) else {
            print("RTMPConnection: 错误 - 无效的RTMP URL: \(rtmpURL)")
            onError?("无效的RTMP URL")
            return false
        }
        
        host = url.host ?? "192.168.1.25"
        port = UInt16(url.port ?? 1935)
        
        print("RTMPConnection: 解析URL成功 - 主机: \(host), 端口: \(port)")
        
        // 测试网络连通性
        testNetworkConnectivity()
        
        // 测试TCP连接
        testTCPConnection()
        
        // 建立TCP连接
        setupConnection()
        
        return true
    }
    
    private func setupConnection() {
        print("RTMPConnection: 开始建立TCP连接...")
        
        let endpoint = NWEndpoint.hostPort(host: NWEndpoint.Host(host), port: NWEndpoint.Port(integerLiteral: port))
        connection = NWConnection(to: endpoint, using: .tcp)
        
        connection?.stateUpdateHandler = { [weak self] state in
            DispatchQueue.main.async {
                self?.handleConnectionState(state)
            }
        }
        
        // 设置连接超时
        connectionTimer = Timer.scheduledTimer(withTimeInterval: connectionTimeout, repeats: false) { [weak self] _ in
            DispatchQueue.main.async {
                print("RTMPConnection: 连接超时")
                self?.onError?("连接超时")
                self?.disconnect()
            }
        }
        
        connection?.start(queue: .global())
    }
    
    private func handleConnectionState(_ state: NWConnection.State) {
        switch state {
        case .ready:
            print("RTMPConnection: TCP连接成功")
            // 取消连接超时定时器
            connectionTimer?.invalidate()
            connectionTimer = nil
            performRTMPHandshake()
        case .failed(let error):
            print("RTMPConnection: TCP连接失败 - \(error)")
            // 取消连接超时定时器
            connectionTimer?.invalidate()
            connectionTimer = nil
            onError?("TCP连接失败: \(error)")
        case .cancelled:
            print("RTMPConnection: TCP连接已取消")
            // 取消连接超时定时器
            connectionTimer?.invalidate()
            connectionTimer = nil
        default:
            break
        }
    }
    
    private func performRTMPHandshake() {
        print("RTMPConnection: 开始RTMP握手...")
        
        // 简化的握手 - 只发送C0
        sendC0()
        
        // 等待S0响应
        receiveS0()
    }
    
    private func sendC0() {
        var data = Data()
        
        // C0: 版本 (1 byte)
        data.append(3) // RTMP version 3
        
        sendData(data)
        print("RTMPConnection: 发送C0完成")
    }
    
    private func receiveS0() {
        receiveData(expectedSize: 1) { [weak self] data in
            guard let data = data, data.count >= 1 else {
                print("RTMPConnection: 接收S0失败")
                self?.onError?("RTMP握手失败")
                return
            }
            
            print("RTMPConnection: 接收S0完成 - 版本: \(data[0])")
            
            // 简化：直接尝试连接
            self?.connectStream()
        }
    }
    
    private func sendC2(_ s1s2Data: Data) {
        // C2: 回显S1的时间戳 + 时间戳 + 随机数据
        var data = Data()
        
        // 回显S1的时间戳 (4 bytes)
        data.append(s1s2Data.prefix(4))
        
        // 当前时间戳 (4 bytes)
        let timestamp = UInt32(Date().timeIntervalSince1970)
        data.append(withUnsafeBytes(of: timestamp.bigEndian) { Data($0) })
        
        // 随机数据 (1528 bytes)
        for _ in 0..<1528 {
            data.append(UInt8.random(in: 0...255))
        }
        
        sendData(data)
        print("RTMPConnection: 发送C2完成")
        
        // 握手完成，开始连接
        connectStream()
    }
    
    private func connectStream() {
        print("RTMPConnection: 开始连接流...")
        
        // 发送connect命令
        sendConnectCommand()
    }
    
    private func sendConnectCommand() {
        let command = createConnectCommand()
        sendRTMPMessage(command, messageType: 0x14, streamId: 0) // INVOKE
        
        // 等待connect响应
        receiveConnectResponse()
    }
    
    private func createConnectCommand() -> Data {
        var data = Data()
        
        // AMF0 字符串 "connect"
        data.append(0x02) // AMF0 string marker
        data.append(0x00) // string length high byte
        data.append(0x07) // string length low byte
        data.append(contentsOf: "connect".utf8)
        
        // AMF0 数字 1.0 (transaction ID)
        data.append(0x00) // AMF0 number marker
        data.append(0x3F) // 0x3F F0 00 00 00 00 00 00 = 1.0
        data.append(0xF0)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        
        // AMF0 对象
        data.append(0x03) // AMF0 object marker
        
        // app
        data.append(0x00) // string length high byte
        data.append(0x03) // string length low byte
        data.append(contentsOf: "app".utf8)
        data.append(0x02) // AMF0 string marker
        data.append(0x00) // string length high byte
        data.append(0x04) // string length low byte
        data.append(contentsOf: "live".utf8)
        
        // flashVer
        data.append(0x00) // string length high byte
        data.append(0x08) // string length low byte
        data.append(contentsOf: "flashVer".utf8)
        data.append(0x02) // AMF0 string marker
        data.append(0x00) // string length high byte
        data.append(0x1F) // string length low byte
        data.append(contentsOf: "FMLE/3.0 (compatible; FMSc/1.0)".utf8)
        
        // swfUrl
        data.append(0x00) // string length high byte
        data.append(0x06) // string length low byte
        data.append(contentsOf: "swfUrl".utf8)
        data.append(0x02) // AMF0 string marker
        data.append(0x00) // string length high byte
        data.append(0x00) // string length low byte (empty string)
        
        // tcUrl
        data.append(0x00) // string length high byte
        data.append(0x05) // string length low byte
        data.append(contentsOf: "tcUrl".utf8)
        data.append(0x02) // AMF0 string marker
        data.append(0x00) // string length high byte
        data.append(0x1A) // string length low byte
        data.append(contentsOf: "rtmp://\(host):\(port)/live".utf8)
        
        // fpad
        data.append(0x00) // string length high byte
        data.append(0x04) // string length low byte
        data.append(contentsOf: "fpad".utf8)
        data.append(0x01) // AMF0 boolean marker
        data.append(0x00) // false
        
        // capabilities
        data.append(0x00) // string length high byte
        data.append(0x0B) // string length low byte
        data.append(contentsOf: "capabilities".utf8)
        data.append(0x00) // AMF0 number marker
        data.append(0x40) // 0x40 3F 00 00 00 00 00 00 = 15.0
        data.append(0x3F)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        
        // audioCodecs
        data.append(0x00) // string length high byte
        data.append(0x0A) // string length low byte
        data.append(contentsOf: "audioCodecs".utf8)
        data.append(0x00) // AMF0 number marker
        data.append(0x40) // 0x40 3F 00 00 00 00 00 00 = 15.0
        data.append(0x3F)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        
        // videoCodecs
        data.append(0x00) // string length high byte
        data.append(0x09) // string length low byte
        data.append(contentsOf: "videoCodecs".utf8)
        data.append(0x00) // AMF0 number marker
        data.append(0x40) // 0x40 3F 00 00 00 00 00 00 = 15.0
        data.append(0x3F)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        
        // videoFunction
        data.append(0x00) // string length high byte
        data.append(0x0C) // string length low byte
        data.append(contentsOf: "videoFunction".utf8)
        data.append(0x01) // AMF0 boolean marker
        data.append(0x01) // true
        
        // objectEncoding
        data.append(0x00) // string length high byte
        data.append(0x0E) // string length low byte
        data.append(contentsOf: "objectEncoding".utf8)
        data.append(0x00) // AMF0 number marker
        data.append(0x3F) // 0x3F F0 00 00 00 00 00 00 = 1.0
        data.append(0xF0)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        
        // 对象结束
        data.append(0x00) // string length high byte
        data.append(0x00) // string length low byte
        data.append(0x09) // AMF0 object end marker
        
        return data
    }
    
    private func receiveConnectResponse() {
        receiveData(expectedSize: 1024) { [weak self] data in
            guard let data = data else {
                print("RTMPConnection: 接收connect响应失败")
                self?.onError?("RTMP连接失败")
                return
            }
            
            print("RTMPConnection: 接收connect响应完成")
            
            // 发送createStream命令
            self?.sendCreateStreamCommand()
        }
    }
    
    private func sendCreateStreamCommand() {
        let command = createCreateStreamCommand()
        sendRTMPMessage(command, messageType: 0x14, streamId: 0) // INVOKE
        
        // 等待createStream响应
        receiveCreateStreamResponse()
    }
    
    private func createCreateStreamCommand() -> Data {
        var data = Data()
        
        // AMF0 字符串 "createStream"
        data.append(0x02) // AMF0 string marker
        data.append(0x00) // string length high byte
        data.append(0x0B) // string length low byte
        data.append(contentsOf: "createStream".utf8)
        
        // AMF0 数字 2.0 (transaction ID)
        data.append(0x00) // AMF0 number marker
        data.append(0x40) // 0x40 00 00 00 00 00 00 00 = 2.0
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        
        return data
    }
    
    private func receiveCreateStreamResponse() {
        receiveData(expectedSize: 1024) { [weak self] data in
            guard let data = data else {
                print("RTMPConnection: 接收createStream响应失败")
                self?.onError?("RTMP流创建失败")
                return
            }
            
            print("RTMPConnection: 接收createStream响应完成")
            
            // 发送publish命令
            self?.sendPublishCommand()
        }
    }
    
    private func sendPublishCommand() {
        let command = createPublishCommand()
        sendRTMPMessage(command, messageType: 0x14, streamId: 1) // INVOKE
        
        // 等待publish响应
        receivePublishResponse()
    }
    
    private func createPublishCommand() -> Data {
        var data = Data()
        
        // AMF0 字符串 "publish"
        data.append(0x02) // AMF0 string marker
        data.append(0x00) // string length high byte
        data.append(0x07) // string length low byte
        data.append(contentsOf: "publish".utf8)
        
        // AMF0 数字 3.0 (transaction ID)
        data.append(0x00) // AMF0 number marker
        data.append(0x40) // 0x40 08 00 00 00 00 00 00 = 3.0
        data.append(0x08)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        data.append(0x00)
        
        // AMF0 null
        data.append(0x05) // AMF0 null marker
        
        // AMF0 字符串 stream name
        data.append(0x02) // AMF0 string marker
        data.append(0x00) // string length high byte
        data.append(UInt8(streamName.count)) // string length low byte
        data.append(contentsOf: streamName.utf8)
        
        // AMF0 字符串 "live"
        data.append(0x02) // AMF0 string marker
        data.append(0x00) // string length high byte
        data.append(0x04) // string length low byte
        data.append(contentsOf: "live".utf8)
        
        return data
    }
    
    private func receivePublishResponse() {
        receiveData(expectedSize: 1024) { [weak self] data in
            guard let data = data else {
                print("RTMPConnection: 接收publish响应失败")
                self?.onError?("RTMP发布失败")
                return
            }
            
            print("RTMPConnection: 接收publish响应完成")
            
            // 连接成功
            DispatchQueue.main.async {
                self?.isConnected = true
                self?.isPublishing = true
                self?.onConnectionStatusChanged?(true)
                self?.onPublishStatusChanged?(true)
                print("RTMPConnection: ✅ RTMP连接和推流成功")
            }
        }
    }
    
    // MARK: - Network Testing
    private func testNetworkConnectivity() {
        print("RTMPConnection: 开始网络连通性测试...")
        
        // 测试SRS API
        let url = URL(string: "http://\(host):1985/api/v1/streams/")!
        
        let task = URLSession.shared.dataTask(with: url) { [weak self] data, response, error in
            DispatchQueue.main.async {
                if let error = error {
                    print("RTMPConnection: SRS API测试失败 - \(error.localizedDescription)")
                } else if let httpResponse = response as? HTTPURLResponse {
                    if httpResponse.statusCode == 200 {
                        print("RTMPConnection: SRS API测试成功 - SRS服务器可达")
                    } else {
                        print("RTMPConnection: SRS API测试失败 - HTTP状态码: \(httpResponse.statusCode)")
                    }
                } else {
                    print("RTMPConnection: SRS API测试失败 - 无效响应")
                }
            }
        }
        task.resume()
    }
    
    private func testTCPConnection() {
        print("RTMPConnection: 开始TCP连接测试...")
        
        let endpoint = NWEndpoint.hostPort(host: NWEndpoint.Host(host), port: NWEndpoint.Port(integerLiteral: port))
        let connection = NWConnection(to: endpoint, using: .tcp)
        
        connection.stateUpdateHandler = { [weak self] state in
            DispatchQueue.main.async {
                self?.handleTCPConnectionState(state)
            }
        }
        
        connection.start(queue: .global())
        
        let startTime = Date()
        let timeout: TimeInterval = 5.0
        
        while Date().timeIntervalSince(startTime) < timeout {
            if connection.state == .ready {
                print("RTMPConnection: TCP连接测试成功")
                connection.cancel()
                return
            }
            Thread.sleep(forTimeInterval: 0.1)
        }
        
        print("RTMPConnection: TCP连接测试失败 - 超时")
        onError?("TCP连接测试失败")
        connection.cancel()
    }
    
    private func handleTCPConnectionState(_ state: NWConnection.State) {
        switch state {
        case .ready:
            print("RTMPConnection: TCP连接测试成功")
        case .failed(let error):
            print("RTMPConnection: TCP连接测试失败 - \(error)")
            onError?("TCP连接测试失败: \(error)")
        case .cancelled:
            print("RTMPConnection: TCP连接已取消")
        default:
            break
        }
    }
    
    // MARK: - Data Sending
    private func sendData(_ data: Data) {
        connection?.send(content: data, completion: .contentProcessed { error in
            if let error = error {
                print("RTMPConnection: 发送数据失败 - \(error)")
            }
        })
    }
    
    private func sendRTMPMessage(_ payload: Data, messageType: UInt8, streamId: UInt32) {
        var header = Data()
        
        print("RTMPConnection: 发送RTMP消息 - 类型: \(messageType), 流ID: \(streamId), 负载大小: \(payload.count)")
        
        // Basic header (1 byte)
        header.append(0x03) // chunk stream ID = 3, fmt = 0
        
        // 修复时间戳溢出问题
        let ms = Date().timeIntervalSince1970 * 1000
        let timestamp = UInt32(ms.truncatingRemainder(dividingBy: Double(UInt32.max)))
        header.append(withUnsafeBytes(of: timestamp.bigEndian) { Data($0) }) // timestamp
        header.append(withUnsafeBytes(of: UInt32(payload.count).bigEndian) { Data($0) }) // message length
        header.append(messageType) // message type
        header.append(0x00) // message stream ID (4 bytes, little endian)
        header.append(0x00)
        header.append(0x00)
        header.append(UInt8(streamId & 0xFF))
        
        print("RTMPConnection: 发送header - 大小: \(header.count) bytes")
        
        // 发送header
        sendData(header)
        
        // 发送payload
        sendData(payload)
        
        print("RTMPConnection: RTMP消息发送完成")
    }
    
    private func receiveData(expectedSize: Int, completion: @escaping (Data?) -> Void) {
        var receivedData = Data()
        let startTime = Date()
        let receiveTimeout: TimeInterval = 10.0
        
        print("RTMPConnection: 开始接收数据 - 期望大小: \(expectedSize) bytes, 超时: \(receiveTimeout)秒")
        
        func receiveMore() {
            // 检查超时
            if Date().timeIntervalSince(startTime) > receiveTimeout {
                print("RTMPConnection: 接收数据超时 - 已接收: \(receivedData.count) bytes")
                completion(nil)
                return
            }
            
            connection?.receive(minimumIncompleteLength: 1, maximumLength: expectedSize) { content, _, isComplete, error in
                if let error = error {
                    print("RTMPConnection: 接收数据失败 - \(error)")
                    completion(nil)
                    return
                }
                
                if let content = content {
                    receivedData.append(content)
                    print("RTMPConnection: 接收到数据 - 大小: \(content.count) bytes, 总计: \(receivedData.count) bytes")
                    
                    if receivedData.count >= expectedSize || isComplete {
                        print("RTMPConnection: 数据接收完成 - 总计: \(receivedData.count) bytes")
                        completion(receivedData)
                    } else {
                        // 继续接收
                        DispatchQueue.global().asyncAfter(deadline: .now() + 0.1) {
                            receiveMore()
                        }
                    }
                } else {
                    print("RTMPConnection: 没有接收到数据")
                    completion(nil)
                }
            }
        }
        
        receiveMore()
    }
    
    func disconnect() {
        print("RTMPConnection: 断开RTMP连接")
        
        // 取消连接超时定时器
        connectionTimer?.invalidate()
        connectionTimer = nil
        
        connection?.cancel()
        connection = nil
        
        isConnected = false
        isPublishing = false
        onConnectionStatusChanged?(false)
        onPublishStatusChanged?(false)
        
        print("RTMPConnection: RTMP连接已断开")
    }
    
    // MARK: - Publishing
    func startPublishing() -> Bool {
        print("RTMPConnection: 开始推流")
        
        // 等待连接完成（最多等待3秒）
        let startTime = Date()
        while !isConnected && Date().timeIntervalSince(startTime) < 3.0 {
            Thread.sleep(forTimeInterval: 0.1)
        }
        
        guard isConnected else {
            onError?("未连接到RTMP服务器")
            return false
        }
        
        return true
    }
    
    func stopPublishing() {
        print("RTMPConnection: 停止推流")
        
        isPublishing = false
        onPublishStatusChanged?(false)
        
        print("RTMPConnection: 推流已停止")
    }
    
    // MARK: - Video Data Sending
    func sendVideoPacket(_ h264Data: Data, timestamp: UInt32) {
        guard isPublishing else { 
            print("RTMPConnection: 跳过视频帧 - 推流未开始")
            return 
        }
        
        // 创建FLV视频标签
        var flvTag = Data()
        
        // FLV Tag Header (11 bytes)
        flvTag.append(0x09) // video tag
        flvTag.append(withUnsafeBytes(of: UInt24(UInt32(h264Data.count)).bigEndian) { Data($0) }) // data size (3 bytes)
        flvTag.append(withUnsafeBytes(of: UInt24(timestamp).bigEndian) { Data($0) }) // timestamp (3 bytes)
        flvTag.append(UInt8((timestamp >> 24) & 0xFF)) // timestamp extended (1 byte)
        flvTag.append(0x00) // stream ID (3 bytes)
        flvTag.append(0x00)
        flvTag.append(0x00)
        
        // FLV Video Tag Body
        flvTag.append(0x17) // video info: keyframe, AVC
        flvTag.append(0x01) // AVC NALU
        flvTag.append(0x00) // composition time (3 bytes)
        flvTag.append(0x00)
        flvTag.append(0x00)
        
        // H.264 data
        flvTag.append(h264Data)
        
        // FLV Tag Footer (4 bytes)
        flvTag.append(withUnsafeBytes(of: UInt32(flvTag.count - 11).bigEndian) { Data($0) }) // previous tag size
        
        // 发送FLV数据
        sendRTMPMessage(flvTag, messageType: 0x09, streamId: 1) // VIDEO
        
        print("RTMPConnection: ✅ 发送视频包成功 - 大小: \(h264Data.count) bytes")
    }
    
    // MARK: - Status
    func getConnectionStatus() -> Bool {
        return isConnected
    }
    
    func getPublishingStatus() -> Bool {
        return isPublishing
    }
}

// 辅助结构体
struct UInt24 {
    let value: UInt32
    
    init(_ value: UInt32) {
        self.value = value & 0xFFFFFF
    }
    
    var bigEndian: UInt32 {
        return value.bigEndian & 0xFFFFFF
    }
} 